package websocket

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
	"tuiwen-go/internal/api"

	"github.com/gofrs/uuid"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		// 允许所有来源，生产环境中请正确配置
		return true
	},
}

// WebSocketMessage 表示来自客户端的传入消息。
type WebSocketMessage struct {
	Type string          `json:"action"`
	Data json.RawMessage `json:"data,omitempty"` // 使用 json.RawMessage 推迟特定处理程序的数据解析
}

// Client 表示一个连接的 WebSocket 客户端。
type Client struct {
	ID   string
	Conn *websocket.Conn
	Send chan interface{} // 用于向客户端发送消息的通道
}

var (
	clients = make(map[string]*Client) // 已连接的客户端
	mu      sync.Mutex                 // 保护客户端映射的互斥锁
)

// HandleConnections 处理传入的 WebSocket 连接
func HandleConnections(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("连接升级失败: %v", err)
		return
	}
	defer ws.Close()

	clientID, _ := uuid.NewV4()
	client := &Client{
		ID:   clientID.String(),
		Conn: ws,
		Send: make(chan interface{}, 256), // Buffered channel for sending messages
	}

	mu.Lock()
	clients[client.ID] = client
	mu.Unlock()

	log.Printf("客户端 %s 已连接", client.ID)
	sendMessage(client, api.WxJSONResponse{Code: 0, Msg: fmt.Sprintf("客户端 %s 已连接", client.ID)})

	// 启动 goroutine 向客户端发送消息
	go client.writePump()

	// 从客户端读取消息
	client.readPump()

	// 客户端断开连接时进行清理
	mu.Lock()
	delete(clients, client.ID)
	close(client.Send)
	mu.Unlock()
	log.Printf("客户端 %s 已断开连接", client.ID)
}

// readPump 将消息从 WebSocket 连接泵送到处理程序。
func (c *Client) readPump() {
	defer func() {
		// 确保客户端从映射中移除并在退出时关闭通道
		mu.Lock()
		delete(clients, c.ID)
		// close(c.Send) // 发送通道由 HandleConnections 在 readPump 返回时关闭
		mu.Unlock()
		c.Conn.Close()
		log.Printf("客户端 %s 连接已关闭, readPump 退出", c.ID)
	}()

	// 如有必要，配置读取限制，例如 c.Conn.SetReadLimit(maxMessageSize)
	// c.Conn.SetReadDeadline(time.Now().Add(pongWait)) // 示例：设置初始读取截止时间
	// c.Conn.SetPongHandler(func(string) error { c.Conn.SetReadDeadline(time.Now().Add(pongWait)); return nil }) // 示例：使用 pong 保持连接

	for {
		messageType, message, err := c.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("客户端 %s 读取消息错误: %v", c.ID, err)
			} else {
				log.Printf("客户端 %s 断开连接 (读取错误): %v", c.ID, err)
			}
			break
		}

		if messageType == websocket.TextMessage {
			log.Printf("从客户端 %s 接收到原始消息: %s", c.ID, string(message))
			var msg WebSocketMessage
			if err := json.Unmarshal(message, &msg); err != nil {
				log.Printf("从客户端 %s 解组消息错误: %v", c.ID, err)
				sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "Invalid message format"})
				continue
			}

			log.Printf("从客户端 %s 解析消息: 类型=%s, 数据=%s", c.ID, msg.Type, string(msg.Data))

			switch msg.Type {
			case "0": // taskNormalHandler
				log.Printf("客户端 %s: 收到消息类型 0 (普通任务)", c.ID)
				taskNormalHandler(c, msg.Data)
			case "1": // taskAddHandler
				log.Printf("客户端 %s: 收到消息类型 1 (添加任务)", c.ID)
				taskAddHandler(c, msg.Data)

			case "generate_scene": // taskAutoSceneHandler
				log.Printf("客户端 %s: 收到消息类型 auto_scene (分镜任务)", c.ID)
				taskAutoSceneHandler(c, msg.Data)
			case "auto_scene": // taskAutoSceneHandler
				log.Printf("客户端 %s: 收到消息类型 auto_scene (分镜任务)", c.ID)
				taskAutoSceneHandler(c, msg.Data)
			case "regenerate-split": // taskTranslateHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-split (单个提示词任务)", c.ID)
				taskAutoSceneHandler(c, msg.Data)
			case "generate_prompt": // taskTranslateHandler
				log.Printf("客户端 %s: 收到消息类型 generate_prompt (提示词任务)", c.ID)
				taskTranslateHandler(c, msg.Data)
			case "regenerate-prompt": // taskTranslateHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-prompt (单个提示词任务)", c.ID)
				taskTranslateHandler(c, msg.Data)
			case "update-prompt": // taskTranslateHandler
				log.Printf("客户端 %s: 收到消息类型 update-prompt (提示词任务)", c.ID)
				taskUpdatePromptHandler(c, msg.Data)
			case "update-split": // taskTranslateHandler
				log.Printf("客户端 %s: 收到消息类型 update-split (分词任务)", c.ID)
				taskUpdateSplitHandler(c, msg.Data)
			case "generate_audio": // taskAudioHandler
				log.Printf("客户端 %s: 收到消息类型 generate_audio (音频任务)", c.ID)
				taskAudioHandler(c, msg.Data)
			case "regenerate-audio": // taskAudioHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-audio (单个音频任务)", c.ID)
				taskAudioHandler(c, msg.Data)
			case "generate_image": // taskImgHandler
				log.Printf("客户端 %s: 收到消息类型 generate_image (图像任务)", c.ID)
				taskImgHandler(c, msg.Data)
			case "regenerate-image": // taskAudioHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-image (单个图像任务)", c.ID)
				taskImgHandler(c, msg.Data)
			case "generate_draft": // taskDraftHandler
				log.Printf("客户端 %s: 收到消息类型 5 (草稿任务)", c.ID)
				taskDraftHandler(c, msg.Data)

			case "extract_roles": // cancelAutoClosePC
				log.Printf("客户端 %s: 收到消息类型 extract_roles (提取角色)", c.ID)
				taskRoleHandler(c, msg.Data)
			case "update-role_ids": // cancelAutoClosePC
				log.Printf("客户端 %s: 收到消息类型 update-role_ids (项目详情挂载角色)", c.ID)
				taskUpdateDetailRoleHandler(c, msg.Data)
			case "regenerate_all": // cancelAutoClosePC
				log.Printf("客户端 %s: 收到消息类型 regenerate_all (一键完成)", c.ID)
				taskRegenerateAllHandler(c, msg.Data)

			case "generate_audio_by_comfyui": // taskImgHandler
				log.Printf("客户端 %s: 收到消息类型 generate_audio_by_comfyui (音频任务)", c.ID)
				taskAudioByComfyuiHandler(c, msg.Data)
			case "regenerate-by-comfyui-audio": // taskAudioHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-by-comfyui-audio (重新生成单个音频任务)", c.ID)
				taskAudioByComfyuiHandler(c, msg.Data)

			case "generate_image_by_comfyui": // taskImgHandler
				log.Printf("客户端 %s: 收到消息类型 generate_image_by_comfyui (图像任务)", c.ID)
				taskImgByComfyuiHandler(c, msg.Data)
			case "regenerate-by-comfyui-image": // taskAudioHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-by-comfyui-image (重新生成单个图像任务)", c.ID)
				taskImgByComfyuiHandler(c, msg.Data)
			case "regenerate-by-comfyui-kontext-imageUp": // taskAudioHandler
				log.Printf("客户端 %s: 收到消息类型 regenerate-by-comfyui-kontext-image (微调图像任务)", c.ID)
				taskImgByComfyuiContextHandler(c, msg.Data)
			case "set_speaker":
				log.Printf("客户端 %s: 收到消息类型 set_speaker (设置讲述人)", c.ID)
				taskSetSpeakerHandler(c, msg.Data)
			default:
				log.Printf("客户端 %s: 收到未知消息类型 %s", c.ID, msg.Type)
				sendMessage(c, api.WxJSONResponse{Code: -1, Msg: fmt.Sprintf("Unknown message type: %s", msg.Type)})
			}
		} else {
			log.Printf("客户端 %s: 收到非文本消息类型: %d", c.ID, messageType)
		}
	}
}

// writePump 将消息从发送通道泵送到 WebSocket 连接。
func (c *Client) writePump() {
	defer func() {
		c.Conn.Close()
		log.Printf("客户端 %s 连接已关闭, writePump 退出", c.ID)
	}()

	for message := range c.Send {
		c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) // 设置写入截止时间
		jsonMessage, err := json.Marshal(message)
		if err != nil {
			log.Printf("客户端 %s 消息编码错误: %v", c.ID, err)
			// 可选：发送错误消息或直接关闭
			return // 如果无法编码消息则退出
		}

		if err := c.Conn.WriteMessage(websocket.TextMessage, jsonMessage); err != nil {
			log.Printf("写入客户端 %s 消息错误: %v", c.ID, err)
			return // 如果写入失败则退出
		}
		log.Printf("发送消息到客户端 %s: %s", c.ID, string(jsonMessage))
	}
}

// sendMessage 是一个辅助函数，用于向客户端发送 JSONResponse。
// 它是非阻塞的，如果发送通道已满则记录日志。
func sendMessage(client *Client, response api.WxJSONResponse) {
	select {
	case client.Send <- response:
	default:
		log.Printf("客户端 %s 发送通道已满, 消息已丢弃: %+v", client.ID, response)
	}
}

// BroadcastMessage 向所有连接的客户端发送消息。
func BroadcastMessage(messageType int, data interface{}) {
	mu.Lock()
	defer mu.Unlock()
	// 根据您的应用程序需求创建消息负载
	// 这是一个通用示例，根据需要进行调整
	payload := api.JSONResponse{
		Code: 0, // 或其他相关代码
		Msg:  fmt.Sprintf("Broadcast message type %d", messageType),
		Data: data,
	}
	for id, client := range clients {
		select {
		case client.Send <- payload:
			log.Printf("广播消息已发送到客户端 %s", id)
		default:
			log.Printf("广播期间客户端 %s 发送通道已满, 消息已丢弃。", id)
		}
	}
}
